###### Namelist for the field initialization of a relativistic electron bunch

import math


dx = 0.4
dtrans = 1.
dt = 0.33
nx =  128 
ntrans = 64 
Lx = nx * dx
Ltrans = ntrans*dtrans
npatch_x = 8 
npatch_r = 8 

# Plasma density
n0 = 0.0035

# Bunch position and rms dimensions (gaussian density distribution)
bunch_sigma_x = 4.
bunch_sigma_r = 3. 
center_bunch = nx*dx/2. 

# Bunch normalized density
alpha = 0.9

# Bunch mean energy
gamma= 200. # relativistic lorentz factor
beta = math.sqrt(1.-1/gamma**2)
relative_energy_spread = 0.01



# normalized density of a bunch with gaussian density
def nbunch_(x,r):
    profile_x = math.exp(-(x-center_bunch)**2/2./bunch_sigma_x**2)
    profile_r = math.exp(-((r)**2)/2./bunch_sigma_r**2)
    profile = alpha*n0*profile_x*profile_r
    
    if (((  (x-center_bunch)**2/(5.*bunch_sigma_x)**2 + ((r)**2)/(5.*bunch_sigma_r)**2 ) < 1. )):
        return profile
    else:
        return 0.


Main(
    geometry = "AMcylindrical",

    interpolation_order = 2,

    timestep = dt,
    simulation_time = 1.*dt, 

    cell_length  = [dx, dtrans],
    grid_length = [ Lx,  Ltrans],

    number_of_AM = 2,

    number_of_AM_relativistic_field_initialization = 1,

    number_of_patches = [npatch_x,npatch_r ],
    
    cluster_width = nx/npatch_x,

    EM_boundary_conditions = [
        ["silver-muller","silver-muller"],
        ["buneman","buneman"],
    ],

    solve_poisson = False,
    
    solve_relativistic_poisson = True,
    relativistic_poisson_max_iteration = 50000,    
    print_every = 100,

)

#MovingWindow(
#    time_start = 0.,
#    velocity_x = 1.
#)

#LoadBalancing(
#    initial_balance = False,
#        every = 20,
#    cell_load = 1.,
#    frozen_particle_load = 0.1
#)


# a more realistic electron bunch should be initialized with numpy arrays
def initial_velocity_longitudinal(x,r):
	return beta

def initial_velocity_radial(x,r):
	return 1e-6

def initial_velocity_theta(x,r):
	return 1e-6

initial_velocity_AM = [initial_velocity_longitudinal,initial_velocity_radial,initial_velocity_theta]

Species(
    name = "bunch_electrons",
    position_initialization = "regular",
    momentum_initialization = "cold",
    relativistic_field_initialization = True,
    particles_per_cell = 20,
    c_part_max = 1.0,
    mass = 1.0,
    charge = -1.0,
    charge_density = nbunch_,
    #mean_velocity = [beta, 0.0, 0.0],
    mean_velocity_AM = initial_velocity_AM, 
    pusher = "boris",
    time_frozen = 0.0,
    boundary_conditions = [
       ["remove", "remove"],
       ["reflective", "remove"],
    ],
)



DiagFields(
    every = 200,
)


DiagProbe(
    every = 200,
    origin   = [0., -ntrans*dtrans,0.],
    corners  = [ [nx*dx,-ntrans*dtrans,0.], [0,ntrans*dtrans,0.] ],
    number   = [nx, 2*ntrans]
)


DiagParticleBinning(
    #name = "my binning",
    deposited_quantity = "weight",
    every = 5,
    time_average = 1,
    species = ["bunch_electrons"],
    axes = [
        ["py",-0.0005, 0.0005, 20],
    ]
)

DiagParticleBinning(
    #name = "my binning",
    deposited_quantity = "weight",
    every = 5,
    time_average = 1,
    species = ["bunch_electrons"],
    axes = [
        ["pz",-0.0005, 0.0005, 20],
    ]
)
